from yolo_resnet34 import YOLOv1_resnet
from dataset import Banana 
from loss import Loss_yolov1
if __name__ == '__main__':
    epoch = 10
    batchsize = 5
    lr = 0.001
    train_data = Banana(is_train=True)
    train_loader = DataLoader(train_data,batch_size=batchsize,shuffle=True)
    net = YOLOv1_resnet()
   optimizer = torch.optim.SGD(net.parameters(),lr=lr,momentum=0.1,weight_decay=0.0005)
    criterion = Loss_yolov1()
    device= torch.device("cpu")
    net = net.to(device)

    start_time = time.time()
    for e in range(epoch):
        net.train()
        for i,(inputs,labels) in enumerate(train_loader):
            # inputs:(5,3,448,448) labels:(5,30,7,7)
            inputs = inputs.to(device)
            labels = labels.float().to(device)
            pred = net(inputs)
            loss = criterion(pred, labels)

            optimizer.zero_grad()  # 清除梯度
            loss.backward()  # 更新梯度
            optimizer.step()  # 更新权重
            print("Epoch %d/%d | Loss:%.2f"%(e,epoch,loss))
    end_time = time.time()
    speed = (end_time - start_time) / 60 / epoch
    print("speed:", speed, "min/epoch")
    torch.save(net.state_dict(),'YOLOV1.pth')  
